//
//  TFCollectionViewSectionInfo.swift
//
//  Created by Andney Jack on 2021/6/20.
//

import UIKit

public struct TFItemBounds {
    public var top: CGFloat
    public var bottom: CGFloat

    public static let zero = TFItemBounds(top: 0, bottom: 0)

    public init(top: CGFloat, bottom: CGFloat) {
        self.top = top
        self.bottom = bottom
    }
}

public class TFCollectionViewSectionInfo: NSObject {
    public var itemRenderDirection: TFCollectionViewItemRenderDirection = .shortestFirst

    public var section: Int = 0
    public var columnCount: Int {
        get {
            return _columnCount
        }
        set {
            _columnCount = newValue
            if _columnHeights != nil {
                _columnHeights?.removeAll()
                for _ in 0 ..< newValue {
                    _columnHeights?.append(0.0)
                }
            }
        }
    }

    public var columnSpacing: CGFloat = 0.0
    public var interitemSpacing: CGFloat = 0.0
    public var collectionViewWidth: CGFloat = 0.0
    public var columnWidth: CGFloat = 0.0
    public var columnWidthClosure: ((_ section: Int, _ columnIndex: Int) -> CGFloat)?

    public var top: CGFloat = 0.0
    public var headerHeight: CGFloat = 0.0
    public var footerHeight: CGFloat = 0.0

    public var headerBounds: TFItemBounds = .zero
    public var footerBounds: TFItemBounds = .zero
    public var headerPinToVisible: Bool = false
    public var footerPinToVisible: Bool = false
    public var hasFloatingHeaderOrFooter: Bool {
        return (headerPinToVisible || footerPinToVisible)
    }

    public var headerInsets: UIEdgeInsets = .zero
    public var footerInsets: UIEdgeInsets = .zero
    public var sectionInsets: UIEdgeInsets = .zero
    public var safeAreaInsets: UIEdgeInsets = .zero

    public var allItemInfos = [TFCollectionViewItemInfo]()
    public var headerAttribute: UICollectionViewLayoutAttributes?
    public var footerAttribute: UICollectionViewLayoutAttributes?

    // MARK: -

    private var _columnCount: Int = 0
    private var _maxColumnHeight: CGFloat = 0.0
    private var _didRecalculate: Bool = false
    private var _columnHeights: [CGFloat]?

    private var columnHeights: [CGFloat] {
        if _columnHeights == nil {
            _columnHeights = Array<CGFloat>(repeating: 0.0, count: columnCount)
        }
        return _columnHeights!
    }

    /// height for header + items + footer
    public var totalSectionHeight: CGFloat {
        return sectionContentHeight + headerTotalHeight + footerTotalHeight
    }

    /// height for items
    public var sectionContentHeight: CGFloat {
        recalculateIfNeeded()
        return _maxColumnHeight
    }

    public var headerTotalHeight: CGFloat {
        return headerInsets.top + headerHeight + headerInsets.bottom
    }

    public var footerTotalHeight: CGFloat {
        return footerInsets.top + footerHeight + footerInsets.bottom
    }

    public var sectionBottom: CGFloat {
        return top + totalSectionHeight
    }

    // MARK: - Public

    public func setNeedRecalculate() {
        _didRecalculate = false
    }

    public func recalculateIfNeeded() {
        if !_didRecalculate {
            m_recalculate()
            _didRecalculate = true
        }
    }

    // MARK: - Private

    private func m_recalculate() {
        if itemRenderDirection == .columnFirst {
            m_itemLayoutByColumnFirst()
        } else {
            m_itemLayoutByWaterfall()
        }

        // header
        if headerHeight > 0 {
            headerAttribute = UICollectionViewLayoutAttributes(forSupplementaryViewOfKind: TFCollectionElementKindSectionHeader,
                                                               with: IndexPath(item: 0, section: section))
            let headerWidth: CGFloat = collectionViewWidth - (safeAreaInsets.left + safeAreaInsets.right) - (headerInsets.left + headerInsets.right)
            let headerFrame = CGRect(x: safeAreaInsets.left + headerInsets.left,
                                     y: top + headerInsets.top,
                                     width: headerWidth,
                                     height: headerHeight)

            headerAttribute?.frame = headerFrame
            headerBounds = TFItemBounds(top: headerFrame.minY, bottom: headerFrame.maxY)
        } else {
            headerAttribute = nil
        }

        // footer
        if footerHeight > 0 {
            footerAttribute = UICollectionViewLayoutAttributes(forSupplementaryViewOfKind: TFCollectionElementKindSectionFooter,
                                                               with: IndexPath(item: 0, section: section))
            let footerWidth: CGFloat = collectionViewWidth - (safeAreaInsets.left + safeAreaInsets.right) - (footerInsets.left + footerInsets.right)
            let footFrame = CGRect(x: safeAreaInsets.left + footerInsets.left,
                                   y: top + headerTotalHeight + _maxColumnHeight,
                                   width: footerWidth,
                                   height: footerHeight)

            footerAttribute?.frame = footFrame
            footerBounds = TFItemBounds(top: footFrame.minY, bottom: footFrame.maxY)
        } else {
            footerAttribute = nil
        }
    }

    // MARK: - Item Layout

    private func m_itemLayoutByWaterfall() {
        let baseTop: CGFloat = top + headerInsets.top + headerHeight + headerInsets.bottom + sectionInsets.top
        let baseOffsetX: CGFloat = safeAreaInsets.left + sectionInsets.left

        for item in 0 ..< allItemInfos.count {
            let itemInfo = allItemInfos[item]
            let columnIndex = nextColumnIndex(for: section, item: item)
            let offsetY = columnHeights[columnIndex]

            itemInfo.columnIndex = columnIndex
            itemInfo.itemAttributes = UICollectionViewLayoutAttributes(forCellWith: IndexPath(item: itemInfo.item, section: itemInfo.section))

            var offsetX = baseOffsetX
            if columnIndex > 0 {
                for _ in 0 ..< columnIndex {
                    offsetX += columnWidth + columnSpacing
                }
            }

            itemInfo.itemAttributes?.frame = CGRect(x: offsetX,
                                                    y: baseTop + offsetY,
                                                    width: columnWidth,
                                                    height: itemInfo.itemHeight)

            _columnHeights?[columnIndex] = offsetY + itemInfo.itemHeight + interitemSpacing
        }

        var maxHeight: CGFloat = 0.0
        for height in columnHeights {
            if height > maxHeight {
                maxHeight = height
            }
        }

        if maxHeight > 0 {
            maxHeight -= interitemSpacing
        }

        _maxColumnHeight = sectionInsets.top + maxHeight + sectionInsets.bottom
    }

    private func m_itemLayoutByColumnFirst() {
        var itemTotalHeight: CGFloat = 0.0
        for item in allItemInfos {
            itemTotalHeight += item.itemHeight + interitemSpacing
        }

        if allItemInfos.count > 0 {
            itemTotalHeight -= interitemSpacing
        }

        let averageHeight = itemTotalHeight / CGFloat(columnCount) // average height for one column
        let baseTop = top + headerInsets.top + headerHeight + headerInsets.bottom + sectionInsets.top

        var offsetX = safeAreaInsets.left + sectionInsets.left
        var offsetY: CGFloat = 0.0
        var columnIndex: Int = 0

        var shouldAddColumnIndex: Bool

        var maxHeight: CGFloat = 0.0
        for idx in 0 ..< allItemInfos.count {
            let itemInfo = allItemInfos[idx]
            shouldAddColumnIndex = false

            if offsetY > 0.0 {
                if allItemInfos.count <= columnCount {
                    shouldAddColumnIndex = true
                } else if columnIndex < columnCount - 1 && offsetY + itemInfo.itemHeight * 0.5 > averageHeight {
                    shouldAddColumnIndex = true
                }
            }

            if shouldAddColumnIndex {
                let lastColumnWidth = columnWidth
                columnIndex += 1

                let previousColumnHeight = offsetY - interitemSpacing
                if maxHeight < previousColumnHeight {
                    maxHeight = previousColumnHeight
                }

                offsetX += lastColumnWidth + columnSpacing
                offsetY = 0.0
            }

            itemInfo.columnIndex = columnIndex
            itemInfo.itemWidth = columnWidth
            itemInfo.itemAttributes = UICollectionViewLayoutAttributes(forCellWith: IndexPath(item: itemInfo.item, section: itemInfo.section))
            itemInfo.itemAttributes?.frame = CGRect(x: offsetX,
                                                    y: offsetY + baseTop,
                                                    width: itemInfo.itemWidth,
                                                    height: itemInfo.itemHeight)

            offsetY += itemInfo.itemHeight + interitemSpacing
        }

        let lastColumnHeight = offsetY > 0.0 ? offsetY - interitemSpacing : 0.0
        if maxHeight < lastColumnHeight {
            maxHeight = lastColumnHeight
        }

        _maxColumnHeight = sectionInsets.top + maxHeight + sectionInsets.bottom
    }

    private func nextColumnIndex(for section: Int, item: Int) -> Int {
        switch itemRenderDirection {
        case .leftToRight:
            return item % columnCount
        case .rightToLeft:
            return (columnCount - 1) - item % columnCount

        default:
            // ShortestFirst
            var columnIndex: Int = 0
            if columnHeights.count > 0 {
                var minHeight: CGFloat = columnHeights[0]
                for idx in 1 ..< columnHeights.count {
                    let height = columnHeights[idx]
                    if height < minHeight {
                        minHeight = height
                        columnIndex = idx
                    }
                }
            }
            return columnIndex
        }
    }
}

public class TFCollectionViewItemInfo: NSObject {
    public var section: Int = 0
    public var item: Int = 0
    public var itemWidth: CGFloat = 0.0
    public var itemHeight: CGFloat = 0.0
    public var columnIndex: Int = 0

    public var itemAttributes: UICollectionViewLayoutAttributes?
}
